import numpy as na
import Image
import sys
import os 
import cPickle
from xml.sax  import make_parser
from CritPoint import CritPoint
import re
import matplotlib.pyplot as pyplot
if( sys.platform == 'darwin' ):
    BIND_TEXT = True
else:
    BIND_TEXT = False
def create(parent):
    return browse(parent)

class browse(object):
    def __init__(self,*args,**kwargs):
        """expecting kwargs to contain some of the following keywords:
            data: a list of arrays to visualize
            mask: a list (possibly empty) of corresponding masks
            grayFactor: a multiplicative scaling for the grayscale
            labelScheme: a dictionary of color,label pairs"""
        try:
            #print kwargs.keys()
            
            self.DEBUG = False
            if( self.DEBUG ): print "browse__init__"
            pointsFile = kwargs.pop('pointsFile',None)
            normalize = kwargs.pop('normalize',False)
            self.axis = kwargs.pop('axis',1)
            scls = kwargs.pop('scls',[1,1,1])
            self.grayFactor = kwargs.pop('grayFactor',[1])
            self.mask = kwargs.pop('mask',[])
            self.scale = 1
            self.absscls = scls
            if( not scls ):
                scls = [1,1,1]
            self.scls = na.array(scls)/min(scls) 
            self.sscl = [self.scls[0],self.scls[1]]
            self.slices = [0,0,0]
            self.plims = [0,1]
            self.view='xy'
            self.wslice = 0
            self.pt = 0
            self.writedir=os.getcwd()
            self.critpoints = []
            self.drawpoints = []
            self.mz = 1
            self.my = 1
            self.mx = 1
            self.mode = "L"
            self.FlipLR = False
            self.FlipUD = False
            self.projections = {'max':na.max,'min':na.min,'average':normSum}
            self.projs = {}
            self.cproj = na.max # current projection
            self.lproj = 'max'
            self.viewtype = 'slice'
            self.win=1
            self.lev=0

            if( pointsFile == None):
                self.pointsFile = "tmp.pckle"
            else:
                self.pointsFile = pointsFile
            if( pointsFile ):
                self.readCriticalPoints()
                
            self.labelScheme = kwargs.pop("labelScheme",{"unlabeled":"Red","labeled":"Green"})
            if( self.labelScheme.has_key("unlabeled") ):
                self.currentLabel = "unlabeled"
            else:
                self.currentLabel = self.labelScheme.keys()[0]

             
            self.rdata = kwargs.pop("data",None)
            if( self.rdata == None ):
                return
            self.filename = os.getcwd()
            self.mergeData()
            if( len(self.rdata[0].shape) != 3 ):
                self.nslices = 1
                self.ProjectSlab()
                vals = getWinLevValues(self.slice)
                self.colorSettings = {'max':vals}
            else:    
                self.getMips()
            self.SetView()
            self.winlevelimage()
        except Exception, error:
            print "failed in browse.__init__() ", error
            sys.exit()

    def getMips(self):
        try:
            self.colorSettings = {}
            if( self.DEBUG ): print "generating projection views"
            views = [('xy',0),('xz',1),('yz',2)]
            for i in range(len(self.data)):
                data = self.data[i]
                for pk in self.projections.keys():
                    cproj = self.projections[pk]
                    for view in views:
                        maxp = self.projs.get(pk,{})
                        tmp = maxp.get(view[0],[])
                        p = cproj(data,axis=view[1])
                        if( not self.colorSettings.has_key(pk) ):
                            vals = getWinLevValues(p)
                            self.colorSettings[pk] = vals
                            
                        tmp.append( p)
                        maxp[view[0]] = tmp
                        self.projs[pk] = maxp

                                   
        except Exception, error:
            print "failed in getMips", error
    def mergeData(self):
        """take self.rawData multiply by mask (if present) and concatenate along specified axis"""
        try:
            self.shape = self.rdata[0].shape
            self.data = []
            for i in range(len(self.rdata)):
                if( self.DEBUG ): print i
                d = self.rdata[i]
                if( self.DEBUG ): print "process mask"
                if( len(self.mask) == 1 ):
                    m = self.mask[0]
                elif( len(self.mask) > 1 ):
                    m = self.mask[i]
                else:
                    m = 1
                if( self.DEBUG ): print "process grayfactor",self.grayFactor
                try:
                    if( len(self.grayFactor) == 1 ):
                        gf = self.grayFactor[0]
                    elif( len(self.grayFactor) > 1 ):
                        gf = self.grayFactor[i]
                except:
                    gf = self.grayFactor
                try:
                    if( m == 1 and gf == 1 ):
                        dm = d
                    else:
                        dm = d * m * f
                except:
                    dm = d
                self.data.append(dm)

            if( len(self.shape) == 3 ):
                self.nslices = self.shape[0]   
                self.mz = self.shape[0]-1
                self.my = self.shape[1]-1
                self.mx = self.shape[2]-1     
            else:
                self.nslices = None
                self.my = self.shape[0]-1
                self.mx = self.shape[1]-1
        except Exception, error:
            print "failed in mergeData", error
    def GetDataFromStandardImage(self, im ):
        try:
            if( im.mode == 'L'):
                self.slice = fromstring(im.tostring(),na.uint8)
                self.slice.shape = im.size[::-1]
                #print "shape when Mono",self.slice.shape
                self.mode = 'L'
            elif( im.mode == 'RGB' ):
                r,g,b = im.split()
                ra = fromstring(r.tostring(),na.uint8)
                ra.shape = r.size[::-1]
                ga = fromstring(g.tostring(),na.uint8)
                ga.shape = g.size[::-1]
                ba = fromstring(b.tostring(),na.uint8)
                ba.shape = b.size[::-1]
                self.slice = na.transpose(na.array([ra,ga,ba]),axes=[1,2,0])
                #print "shape when RGB",self.slice.shape[0:2]
                self.mode = 'RGB'
            else:
                self.mode = None
                self.slice = None
        except Exception, error:
            print "failed in GetDataFromStandardImage() ", error

    def win_lev(self, data, maxc = 255 ):
        try:
            #print "self.lev",self.lev, self.win
            win = self.win
            lev = self.lev
            m = maxc / (2.0*win)
            o = m*(lev-win)
            temp= na.clip((m*data-o),0,maxc).astype(na.uint8)
            #print "win_lev result range",temp.max(),temp.min()
            return temp
        except Exception, error:
            print "failed in win_lev() ", error
    def winlevelimage(self):
        try:
            if( self.DEBUG ): print self.slice.max()
            temp = self.win_lev(self.slice)
            # why is my slice coming out upside down?
            temp = na.transpose(na.array((temp,temp,temp),na.uint8),(1,2,0))
            self.rawbuffer = na.require(temp,na.uint8,'CO')
        except Exception, error:
            print "failed in winlevelimage", error


    def ProjectSlab(self):
        try:
            if( self.nslices > 1 ):
                if( self.DEBUG ): print "ProjectSlab"
                self.dataSlices = []
                if( self.view == 'xy' ):
                    self.slices[0] = self.wslice
                    self.plims = [max(0,self.wslice-self.pt),
                                  min(self.shape[0],self.wslice+self.pt+1)]
                    for d in self.data:
                        self.dataSlices.append(self.cproj(d[self.plims[0]:self.plims[1],:,:],0))
                elif( self.view == 'xz' ):
                    self.slices[1] = self.wslice
                    self.plims = [max(0,self.wslice-self.pt),min(self.shape[1],self.wslice+self.pt+1)]
                    for d in self.data:
                        self.dataSlices.append(self.cproj(d[:,self.plims[0]:self.plims[1],:],1))
                else:
                    self.slices[2] = self.wslice
                    self.plims = [max(0,self.wslice-self.pt),min(self.shape[2],self.wslice+self.pt+1)]
                    for d in self.data:
                        self.dataSlices.append(self.cproj(d[:,:,self.plims[0]:self.plims[1]],2))
            else:
                self.dataSlices = []
                for d in self.data:
                    self.dataSlices.append(d)

            self.makeSlice()
        except Exception, error:
            print "failed in ProjectSlab", error
    def makeSlice(self):
        self.slice = na.concatenate(self.dataSlices,axis=self.axis)
        self.slice = self.slice[::-1,:]
        if(self.FlipLR):
            if( self.mode == 'L' ):
                self.slice = self.slice[::-1,:]
            elif( self.mode == 'RGB' ):
                self.slice = self.slice[:,::-1,:]
        if(self.FlipUD):
            if( self.mode == 'L' ):
                self.slice = self.slice[:,::-1]
            elif( self.mode == 'RGB' ):
                self.slice = self.slice[:,:,::-1]
        if( self.DEBUG ): print "makeSlice",self.slice.shape,self.slice.max()
    def writeCriticalPoints(self):
        try:
            fo = open(self.pointsFile,"wb")
            cPickle.dump(self.critpoints,fo)
            fo.close()
        except Exception, error:
            print "failed in writeCriticalPoints", error
    def readCriticalPoints(self):
        try:
            self.writedir = os.path.dirname(self.pointsFile)
            fo = open(self.pointsFile,"rb")
            newPoints = cPickle.load(fo)
            
            self.critpoints.extend(newPoints)
            fo.close()
        except Exception, error:
            print "failed in readCriticalPoints", error
    def findNearestCritPoint(self,p,verbose=True):
        try:
            minind = 0 
            minpoint = self.critpoints[0]
            mindist = (minpoint[0]-p[0])**2+(minpoint[1]-p[1])**2+(minpoint[2]-p[2])**2
            for i in range(1,len(self.critpoints)):
                point = self.critpoints[i]
                dist = (point[0]-p[0])**2+(point[1]-p[1])**2+(point[2]-p[2])**2 
                if( dist < mindist ):
                    mindist = dist
                    minind = i
                    minpoint = point
            if( verbose ):
                print minind, minpoint, mindist
            return minind, minpoint, mindist
        except Exception, error:
            print "failed in findNearestCritPoint", error
            return None

    def delete_nearest_critpoint(self, p ):
        try:
            minind, minpoint, mindist = self.findNearestCritPoint(p)
            if( minind != None ):
                del self.critpoints[minind]
            return minpoint
        except Exception, error:
            print "failed in delete nearest_drawpointXY",error
    def modifyNearestCritPoint(self,p):
        minind, minpoint, mindist = self.findNearestCritPoint(p)
        self.critpoints[minind].setLabel(self.currentLabel)
        self.DrawCritPoints()

    def DrawAllCritPoints(self):
        try:
            self.drawpoints = self.critpoints
        except:
            pass

    def DrawCritPointsXY(self):
        try:
            #print "DrawCritPointsXY"
            self.drawpoints = []
            for point in self.critpoints:
                if( self.plims[0] <= point[2] < self.plims[1] ):
                    self.drawpoints.append(point)
        except:
            pass

    def DrawCritPointsXZ(self):
        try:
            if( self.DEBUG ): print "DrawCritPointsXZ"
            self.drawpoints = []
            for point in self.critpoints:
                if( self.plims[0] <= point[1] < self.plims[1] ):
                    self.drawpoints.append(point)
        except:
            pass
      
    def DrawCritPointsYZ(self):
        try:
            if( self.DEBUG ): print "DrawCritPointsYZ"
            self.drawpoints = []
            for point in self.critpoints:
                if( self.plims[0] <= point[0] < self.plims[1] ):
                    self.drawpoints.append(point)
        except:
            pass
    def DrawCritPoints(self):
        try:
            if( self.DEBUG ): print "DrawCritPoints"
            if( self.view == 'xy' ):
                self.DrawCritPointsXY()
            elif( self.view == 'xz' ):
                self.DrawCritPointsXZ()
            else:
                self.DrawCritPointsYZ()
        except Exception, error:
            print "failed in DrawCritPoints", error

    def invertSlice(self):
        try:
            if( self.mode == 'L'):
                self.slice = self.slice.max()-self.slice
            elif( self.mode == 'RGB' ):
                self.slice[:,:,0] = self.slice[:,:,0].max()-self.slice[:,:,0]
                self.slice[:,:,1] = self.slice[:,:,1].max()-self.slice[:,:,1]
                self.slice[:,:,2] = self.slice[:,:,2].max()-self.slice[:,:,2]
            else:
                return
            self.winlevelimage()
        except:
            pass
    def ResetScale(self):
        try:
            self.scale = 1
        except:
            pass
    def DecreaseScale(self):
        try:
            self.scale = max(1,self.scale-1)
        except:
            pass
    def IncreaseScale(self):
        try:
            self.scale = self.scale+1
        except:
            pass
    def SetView(self):
        try:
            self.cproj = self.projections[self.lproj]
            if( self.viewtype == 'projection' ):
                mips = self.projs[self.lproj]
                self.dataSlices = []
                for m in mips[self.view]:
                    self.dataSlices.append(m)
        
                if( self.view == 'xy' ):
                    self.plims=[0,self.shape[0]]
                elif( self.view == 'xz' ):
                    self.plims = [0,self.shape[1]]
                else:
                    self.plims = [0,self.shape[2]]
                self.makeSlice()
                self.DrawAllCritPoints()
            else:
                self.ProjectSlab()
                self.DrawCritPoints()
        except Exception, error:
            print "failed in SetView", error

    def SetViewXY(self):
        try:
            self.view = 'xy'
            self.sscl = [self.scls[0],self.scls[1]]
            self.wslice = self.slices[0]
            self.SetView()
            self.winlevelimage()
        except Exception, error:
            pass
    def SetViewXZ(self):
        try:
            self.view = 'xz'
            self.sscl = [self.scls[0],self.scls[2]]
            self.wslice = self.slices[1]
            self.SetView()
            self.winlevelimage()
        except Exception, error:
            print "failed in SetViewXZ",error
            pass
    def SetViewYZ(self):
        try:
            self.view = 'yz'
            self.sscl = [self.scls[1],self.scls[2]]
            self.wslice = self.slices[2]
            self.SetView()
            self.winlevelimage()
        except Exception, error:
            print "failed in SetViewYZ", error


    def flipSliceUD(self):
        try:
            self.FlipUD = not self.FlipUD
            self.makeSlice()
            self.winlevelimage()
        except:
            pass

    def flipSliceLR(self):
        try:
            self.FlipLR = not self.FlipLR
            self.makeSlice()
            self.winlevelimage()
        except:
            pass

    def transposeSlice(self):
        try:
            self.slice = na.transpose(self.slice)
        except:
            pass
def getLabels(filename):
    try:
        fo = open(filename,'r')
        labels = fo.readlines()
        labels = [ l.strip() for l in labels ]
        return labels
    except Exception, error:
        print "failed in getLabelScheme", error
def getColors(filename):
    try:
        fo = open(filename,'r')
        colors = fo.readlines()
        colors = [ c.strip() for c in colors ]
        return colors
    except Exception, error:
        print "failed in getColors", error
def getLabelScheme(labels, colors):
    try:
        colorScheme = {}
        # always use the first color uniquely for the first label
        colorScheme[labels[0]] = colors[0]

        numColors = len(colors)-1
        count = 0
        for l in labels[1:]:
            colorScheme[l] = colors[count % numColors +1]
            count += 1
        return colorScheme
    except Exception, error:
        print "failed in getLabelScheme", error

def normSum(data,axis=0):
    mx = data.max()
    d = na.sum(data-data.min(),axis=axis)
    temp= 10*na.log(na.abs(d)+1)
    #print temp.min(),temp.max()
    return temp
def getLabels(filename):
    try:
        fo = open(filename,'r')
        labels = fo.readlines()
        labels = [ l.strip() for l in labels ]
        return labels
    except Exception, error:
        print "failed in getLabelScheme", error
def getWinLevValues(data):
    imax = data.max()
    imin = data.min()
    wh = (imax-imin)
    wl = 1
    lh = (imax+imin)/2+(imax-imin)-1
    ll = (imax+imin)/2-(imax-imin)
    vals = {'ll':ll,'lh':lh,'wh':wh,'wl':wl,'wc':(wl+wh/3),'lc':(ll+lh)/2}
    return vals


